#!/usr/bin/env python
# coding: utf-8

"""
labelme2yolo.py
create by cuiyu 2025-02-19
说明：
  * 将 labelme 的 json 标注文件转换为 yolo 格式的 txt 文件
  * 标签和 ID 的映射关系从 labelme2yolo-rddhw.json 中读取，格式为：(labelx 不包含前导数组)
    *   id1: label1
    *   id2: label2

用法:
  labelme2yolo.py

更新:
2025-02-19:
"""

import json 
from pathlib import Path
from loguru import logger
import datetime
import sys
from bhclasses_define import BHClasses

_APP_FILE = Path(sys.argv[0])
_APP_ROOT = _APP_FILE.parent
_APP_NAME = _APP_FILE.stem

class Labelme2Yolo:
    def __init__(self, img_root: Path):
        """
        将 labelme 的 json 标注文件转换为 yolo 格式的 txt 文件
        txt 文件位于 json 文件所在的目录下，文件名与 json 文件名相同，后缀为 txt

        Args:
          :img_root: Path,  输入目录

        Returns: None
        """
        self.img_root = Path(img_root)

        # 存储从所有标注文件中找到的 classes, key=id, val=name
        # self.classes_model: dict={
        #     "0": "a车道线",
        #     "1": "1横向裂缝",
        #     "2": "2纵向裂缝",
        #     "3": "3龟裂",
        #     "4": "4横向条修",
        #     "5": "5纵向条修",
        #     "6": "6块修",
        #     "7": "7松散",
        #     "8": "8横向条修不良",
        #     "9": "9纵向条修不良",
        #     "10": "c施工缝",
        #     "11": "b伸缩缝"
        # }      

        # 存储从所有标注文件中找到的 classes, key=name, val=id
        # self.classes_labelme: dict = dict(zip(self.classes_model.values(), self.classes_model.keys()))  # 需要交换 Key Val

        # classes 存储 labelme 标签和 模型类型 ID 的映射关系， 从 self.classes_file 文件中加载
        # _classes: dict = json.loads(self.classes_file.read_text(encoding="utf-8"))
        # self.classes = dict(zip(_classes.values(), _classes.keys()))  # 需要交换 Key Val
        pass


    def json2yolo(self, json_file: Path) -> bool:
        """
        将 json_file 文件转换为 yolo 格式的 txt 文件
        txt 文件位于 json 文件所在的目录下，文件名与 json 文件名相同，后缀为 txt

        Args:
            :json_file: Path,  json 文件路径

        Returns: bool, 是否转换成功
        """
        json_str = json_file.read_text(encoding="utf-8")
        try:
            json_dict = json.loads(json_str)
        except:
            logger.error(f"解析 json 文件失败 {str(json_file)}")
            return False

        h,w = json_dict['imageHeight'], json_dict['imageWidth']
        txt_path = json_file.with_suffix(".txt")

        with open(txt_path, "w") as f:
            for shape_dict in json_dict['shapes']:
                # 获取标签
                label = shape_dict['label']
                # label = label[1:]   # 删除第一个字符

                # 处理标签列表
                label_index = BHClasses.getID(label)
                
                points = shape_dict['points']
                points_nor_list=[]
                for point in points:
                    points_nor_list.append(point[0]/w)
                    points_nor_list.append(point[1]/h)
                points_nor_list=list(map(lambda x:str(x),points_nor_list))
                points_nor_str=' '.join(points_nor_list)
                label_str=str(label_index)+' '+points_nor_str+'\n'
                f.writelines(label_str)

        return True
        

    def labelme2yolo(self):
        """
        将 labelme 的 json 标注文件转换为 yolo 格式的 txt 文件  
        txt 文件位于 json 文件所在的目录下，文件名与 json 文件名相同，后缀为 txt  
        """

        # 2. 按照文件名排序返回 img_root 下的 .json 文件列表，准备进度窗
        json_files = [Path(file) for file in sorted(self.img_root.rglob("*.json"), key=lambda file: str(file))]

        # 3. 循环处理每个 json 文件，并显示进度
        num = len(json_files)
        n_len = len(str(num))
        for i, file in enumerate(json_files):
            n = i+1
            print(f'{n:{n_len}} / {num} {float(n)/num*100:.2f}%  {str(file.name)}', end='\r')
            self.json2yolo(file)
        print("")
        pass

    @staticmethod        
    def json_delimgs(img_root: Path):
        """
        将 labelme 标注文件 .json 中嵌入的图片删除
        
        Args:
            :img_root: Path,  输入目录  
        
        Returns: None
        """
        def _json_delimg(json_file: Path) -> bool:
            """
            将 json 文件中嵌入的图片删除
            Returns: bool, 是否转换成功
            """
            json_str = json_file.read_text(encoding="utf-8")
            try:
                json_dict = json.loads(json_str)
            except:
                logger.error(f"解析 json 文件失败 {str(json_file)}")
                return False

            json_dict["imageData"] = None

            json_str = json.dumps(json_dict, ensure_ascii=False, indent=4)
            json_file.write_text(json_str, encoding="utf-8")

            return True

        # 1.
        img_root = Path(img_root)

        # 2. 按照文件名排序返回 img_root 下的 .json 文件列表，准备进度窗
        json_files = [Path(file) for file in sorted(img_root.rglob("*.json"), key=lambda file: str(file))]

        # 3. 循环处理每个 json 文件，并显示进度
        num = len(json_files)
        n_len = len(str(num))
        for i, file in enumerate(json_files):
            n = i+1
            print(f'{n:{n_len}} / {num} {float(n)/num*100:.2f}%  {str(file.name)}', end='\r')
            _json_delimg(file)
        print("")
        pass
   
            

def _logInit() -> None:
    # logger.add(sys.stderr, format="{time} {level} {message}", level="INFO")
    time = datetime.datetime.today()
    # f'the time is {time:%Y-%m-%d (%a) %H:%M:%S}'   # datetime时间格式
    logger.add(f"{_APP_ROOT}/log/{_APP_NAME}-{time:%Y%m%d}.log", rotation="500 MB", encoding='utf-8', retention="10 days")  
    logger.add(f"{_APP_ROOT}/log/{_APP_NAME}-err-{time:%Y%m%d}.log", rotation="500 MB", encoding='utf-8', level='ERROR', retention="30 days")  

    # t = time.strftime("%Y%m%d") 
    # logger.add(f"{baseDir}/log-{t}.log", rotation="500 MB", encoding='utf-8', retention="10 days")  
    # logger.add(f"{baseDir}/err-{t}.log", rotation="500 MB", encoding='utf-8', level='ERROR', retention="30 days")  

def _main():
    img_root = Path("/aaby/netdisk/mydata/ai-data/dataset/priv/rddhw/训练数据/img01")
    # classes_file_id = _APP_ROOT / "classes-rddhw-id.json"
    # classes_file_labelme = Path(_APP_ROOT)/classes_labelme.json
    # l2y = Labelme2Yolo(img_root, classes_file_id, classes_file_labelme
    l2y = Labelme2Yolo(img_root)

    # 将 labelme 的 json 标注文件转换为 yolo 格式的 txt 文件
    print("1. 将 labelme 的 json 标注文件转换为 yolo 格式的 txt 文件：")
    l2y.labelme2yolo()
    print(l2y.classes_labelme)
    print(l2y.classes_model)
    print(".")

    # 将 json 标注文件中嵌入的图片删除
    # print("2. 将 labelme 标注文件 .json 中嵌入的图片删除：")
    # l2y.json_delimgs(img_root)
    # print(".")

if __name__=="__main__":
    _logInit()
    _main()
